import faiss
from typing import Union, List, Tuple

import numpy


class BasicIndexerHelper:
    def __init__(self):
        pass

    def update(self, *args, **kw):
        raise NotImplementedError

    def search(self, *args, **kw):
        raise NotImplementedError

    def save(self, *args, **kw):
        raise NotImplementedError

    def load(self, *args, **kw):
        raise NotImplementedError

    def clear(self):
        raise NotImplementedError

    def getindexer(self):
        raise NotImplementedError


class FaissIndexerHelper(BasicIndexerHelper):
    def __init__(self, dim=50, back='indexflatl2'):
        """
        faiss的索引帮助器，用于提供简易的索引管理对象。
        :param dim: 储存的索引维度。
        :param back: 使用的Faiss索引对象。
        """
        super().__init__()
        if back == 'indexflatl2':
            self.indexer = faiss.IndexFlatL2(dim)
        else:
            raise Exception("不支持的Indexer类型。")

    def update(self, data):
        """
        向索引器中添加数据。
        :param data:
        :return:
        """
        self.indexer.add(data)

    def search(self, query: Union[str, List[str]], k=1) -> Tuple[numpy.ndarray, List[int]]:
        """
        查询向量`query`最接近的向量。
        :param query: 查询的向量。
        :param k: 返回的条数。
        :return: 返回距离、查询结果向量的索引。
        """
        D: numpy.ndarray
        I: List[int]
        D, I = self.indexer.search(query, k)
        return D, I

    def save(self, path):
        """
        保存向量数据库。
        :param path:
        :return:
        """
        faiss.write_index(self.indexer, path)

    def load(self, path):
        """
        导入向量数据库。
        :param path:
        :return:
        """
        self.indexer = faiss.read_index(path)

    def clear(self):
        self.indexer.reset()

    def getindexer(self):
        return self.indexer
